import os
import torch
import sys
import torch.nn as nn
from torch.utils.data import DataLoader
from dataset.Data_Loader import LOAD_DATA
from tqdm import tqdm
from model.PResnet_model import PResnet
from params import params_dict
from utils.utils import *
import random
from utils.evaluate import *

def fit(model, epoch, train_loader, valid_loader, test_loader, val_num, optimizer, loss_function, params_dict):
    model.train()
    running_loss = 0.0
    train_bar = tqdm(train_loader, file=sys.stdout)
    for step, data in enumerate(train_bar):
        imgs, features_seqs, labels, _ = data
        optimizer.zero_grad()
        outputs = model(imgs.to(params_dict['device']), features_seqs.to(params_dict['device']))
        loss = loss_function(outputs, labels.to(params_dict['device']))
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

        train_bar.desc = "imgs epoch[{}/{}] loss:{:.3f}".format(epoch + 1, params_dict['epoch'], loss)
    model.eval()
    acc = 0.0
    with torch.no_grad():
        val_bar = tqdm(valid_loader, file=sys.stdout)
        for val_data in val_bar:
            val_imgs, val_features_seqs, val_labels, _ = val_data
            outputs = model(val_imgs.to(params_dict['device']), val_features_seqs.to(params_dict['device']))
            predict_y = torch.max(outputs, dim=1)[1]
            acc += torch.eq(predict_y, val_labels.to(params_dict['device'])).sum().item()
    class_list = []
    label_list = []
    with torch.no_grad():
        test_bar = tqdm(test_loader, file=sys.stdout)
        for test_data in test_bar:
            test_imgs, test_features_seqs, test_labels, _ = test_data
            outputs = model(test_imgs.to(params_dict['device']), test_features_seqs.to(params_dict['device']))
            pred = torch.softmax(outputs, dim=1)
            predict_cla = torch.argmax(pred, dim=1)
            class_list.extend(predict_cla.to('cpu').numpy())
            label_list.extend(list(map(int, test_labels)))
    pred_tensor = torch.tensor(class_list)
    true_tensor = torch.tensor(label_list)
    F1_score = test_f1_en(pred_tensor, true_tensor)
    test_accurate = test_acc_en(pred_tensor, true_tensor)
    Specificity = test_spc_en(pred_tensor, true_tensor)
    recall = test_rcl_en(pred_tensor, true_tensor)
    Precision = test_pcl_en(pred_tensor, true_tensor)
    Kapa = cohen_kappa_score(pred_tensor, true_tensor)

    val_accurate = acc / val_num
    return running_loss, val_accurate, test_accurate, Kapa, F1_score, Specificity, recall, Precision


def main(params_dict):
    train_dataset = LOAD_DATA(params_dict['train_img_path'], params_dict['train_file_path'], mode='imgs')
    val_dataset = LOAD_DATA(params_dict['valid_img_path'], params_dict['valid_file_path'], mode='valid')
    test_dataset = LOAD_DATA(params_dict['test_img_path'], params_dict['test_file_path'], mode='test')
    nw = params_dict['nw']
    val_num = len(val_dataset)
    train_loader = DataLoader(train_dataset, batch_size=params_dict['batch_size'], shuffle=True, num_workers=nw)
    valid_loader = DataLoader(val_dataset, batch_size=params_dict['batch_size'], shuffle=False, num_workers=nw)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=nw)
    model = load_model(params_dict)
    model.to(params_dict['device'])
    optimizer = set_optimizer(model, params_dict['lr'])
    loss_function = set_loss()
    best_model = 0.0
    train_steps = len(train_loader)
    for epoch in range(params_dict['epoch']):
        running_loss, val_accurate, test_accurate, Kapa, F1_score, Specificity, recall, Precision \
        = fit(model, epoch, train_loader, valid_loader, test_loader, val_num, optimizer, loss_function, params_dict)
        print('[epoch %d] train_loss: %.3f  val_accuracy: %.3f test_accurate: %.3f Kapa: %.3f F1_score: %.3f' %
              (epoch + 1, running_loss / train_steps, val_accurate, test_accurate, Kapa, F1_score))
        save_model(model, params_dict['save_file_path'], params_dict['save_last_model_path'])
        if val_accuracy > best_model:
            best_model = val_accuracy
            save_model(model, params_dict['save_file_path'], params_dict['save_best_model_path'])
    print("Training finsh!!!")


if __name__ == '__main__':
    main(params_dict)
